package crypto

import (
	"crypto/aes"
	"encoding/hex"
	"github.com/pkg/errors"
)

type AesKey struct {
	Length int
	Data   []byte
}

func (k *AesKey) GetData() []byte {
	return k.Data
}
func (k *AesKey) StoreFromHex(data string) error {
	result, err := hex.DecodeString(data)
	if err == nil {
		copy(k.Data, result)
	}
	return err
}

func (k *AesKey) IsZero() bool {
	return k.Data[0]|k.Data[1] == 0
}

func NewDefaultAesKey() *AesKey {
	return NewAesKey(0x10)
}

func NewAesKey(length int) *AesKey {
	return &AesKey{
		Length: length,
		Data:   make([]byte, length),
	}
}

func NewAesKeyWithData(length int, data []byte) *AesKey {
	return &AesKey{
		length, data,
	}
}

type AccessKey struct {
	*AesKey
}

func NewAccessKey(titleKey []byte) *AccessKey {
	return &AccessKey{
		AesKey: NewAesKeyWithData(128, titleKey),
	}
}

type AesXtsKey struct {
	*AesKey
}

func (k *AesXtsKey) SubKeys(i int) []byte {
	data := make([]byte, 0x10)
	copy(data, k.Data[0x10*i:])
	return data
}

func NewAesXtsKey() *AesXtsKey {
	return &AesXtsKey{
		AesKey: NewAesKey(0x20),
	}
}

type EncryptedKeyBlob struct {
	*AesKey
}

func NewEncryptedKeyBlob() *EncryptedKeyBlob {
	return &EncryptedKeyBlob{
		AesKey: NewAesKey(144),
	}
}

type KeyBlob struct {
	*AesKey
}

func NewKeyBlob() *KeyBlob {
	return &KeyBlob{
		AesKey: NewAesKey(144),
	}
}

type TransformFunc func(dst, src []byte)

func NewAesECBEncrypt(key []byte) (TransformFunc, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, errors.WithStack(err)
	}
	return func(encrypted, origData []byte) {
		size := block.BlockSize()
		for bs, be := 0, size; bs < len(origData); bs, be = bs+size, be+size {
			block.Encrypt(encrypted[bs:be], origData[bs:be])
		}
	}, nil
}

func NewAesECBDecrypt(key []byte) (TransformFunc, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, errors.WithStack(err)
	}
	return func(decrypted, encrypted []byte) {
		size := block.BlockSize()
		for bs, be := 0, size; bs < len(encrypted); bs, be = bs+size, be+size {
			block.Decrypt(decrypted[bs:be], encrypted[bs:be])
		}
	}, nil
}

func AesECBDecrypt(input, key []byte) ([]byte, error) {
	decrypt, err := NewAesECBDecrypt(key)
	if err != nil {
		return nil, err
	}
	bytes := make([]byte, len(input))
	decrypt(bytes, input)
	return bytes, nil
}
